import torch
import torch.nn as nn
from torchvision import transforms, datasets
from torch.utils.data import DataLoader
from torchsummary import summary
from torchviz import make_dot
import matplotlib.pyplot as plt
import time
import csv
import os
 
class BasicBlock(nn.Module):
    expansion = 1
 
    def __init__(self, in_channel, out_channel, stride=1, downsample=None):
        super(BasicBlock, self).__init__()
 
        self.conv1 = nn.Conv2d(in_channel, out_channel, kernel_size=3, padding=1, stride=stride, bias=False)
        self.bn1 = nn.BatchNorm2d(out_channel)    # BN层, BN层放在conv层和relu层中间使用
        self.conv2 = nn.Conv2d(out_channel, out_channel, kernel_size=3, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_channel)

        self.downsample = downsample
        self.relu = nn.ReLU(inplace=True)
 
    def forward(self, X):
        identity = X
        Y = self.relu(self.bn1(self.conv1(X)))
        Y = self.bn2(self.conv2(Y))
 
        if self.downsample is not None:    # 保证原始输入X的size与主分支卷积后的输出size叠加时维度相同
            identity = self.downsample(X)
 
        return self.relu(Y + identity)

 
class BottleNeck(nn.Module):
    # BottleNeck模块最终输出out_channel是Residual模块输入in_channel的size的4倍(Residual模块输入为64)，shortcut分支in_channel
    # 为Residual的输入64，因此需要在shortcut分支上将Residual模块的in_channel扩张4倍，使之与原始输入图片X的size一致
    expansion = 4

    def __init__(self, in_channel, out_channel, stride=1, downsample=None):
        super(BottleNeck, self).__init__()
        # 默认原始输入为224，经过7x7层和3x3层之后BottleNeck的输入降至64
        self.conv1 = nn.Conv2d(in_channel, out_channel, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(out_channel)    # BN层, BN层放在conv层和relu层中间使用
        self.conv2 = nn.Conv2d(out_channel, out_channel, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_channel)
        self.conv3 = nn.Conv2d(out_channel, out_channel * self.expansion, kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm2d(out_channel * self.expansion)  # Residual中第三层out_channel扩张到in_channel的4倍
 
        self.downsample = downsample
        self.relu = nn.ReLU(inplace=True)
 
    # 前向传播
    def forward(self, X):
        identity = X
 
        Y = self.relu(self.bn1(self.conv1(X)))
        Y = self.relu(self.bn2(self.conv2(Y)))
        Y = self.bn3(self.conv3(Y))
 
        if self.downsample is not None:    # 保证原始输入X的size与主分支卷积后的输出size叠加时维度相同
            identity = self.downsample(X)
 
        return self.relu(Y + identity)
 
 
class ResNet(nn.Module):
    # num_classes是训练集的分类个数，include_top是在ResNet的基础上搭建更加复杂的网络时用到，此处用不到
    def __init__(self, residual, num_residuals, num_classes=10, include_top=True):
        super(ResNet, self).__init__()
 
        self.out_channel = 64    # 输出通道数(即卷积核个数)，会生成与设定的输出通道数相同的卷积核个数
        self.include_top = include_top
 
        self.conv1 = nn.Conv2d(3, self.out_channel, kernel_size=7, stride=2, padding=3,
                               bias=False)    # 3表示输入特征图像的RGB通道数为3，即图片数据的输入通道为3
        self.bn1 = nn.BatchNorm2d(self.out_channel)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        self.conv2 = self.residual_block(residual, 64, num_residuals[0])
        self.conv3 = self.residual_block(residual, 128, num_residuals[1], stride=2)
        self.conv4 = self.residual_block(residual, 256, num_residuals[2], stride=2)
        self.conv5 = self.residual_block(residual, 512, num_residuals[3], stride=2)
        if self.include_top:
            self.avgpool = nn.AdaptiveAvgPool2d((1, 1))    # output_size = (1, 1)
            # self.dropout = nn.Dropout(0.5)
            self.fc = nn.Linear(512 * residual.expansion, num_classes)
 
        # 对conv层进行初始化操作
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
 
    def residual_block(self, residual, channel, num_residuals, stride=1):
        downsample = None
 
        # 用在每个conv_x组块的第一层的shortcut分支上，此时上个conv_x输出out_channel与本conv_x所要求的输入in_channel通道数不同，
        # 所以用downsample调整进行升维，使输出out_channel调整到本conv_x后续处理所要求的维度。
        # 同时stride=2进行下采样减小尺寸size，(注：conv2时没有进行下采样，conv3-5进行下采样，size=56、28、14、7)。
        if stride != 1 or self.out_channel != channel * residual.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(self.out_channel, channel * residual.expansion, kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(channel * residual.expansion))
 
        block = []    # block列表保存某个conv_x组块里for循环生成的所有层
        # 添加每一个conv_x组块里的第一层，第一层决定此组块是否需要下采样(后续层不需要)
        block.append(residual(self.out_channel, channel, downsample=downsample, stride=stride))
        self.out_channel = channel * residual.expansion    # 输出通道out_channel扩张
 
        for _ in range(1, num_residuals):
            block.append(residual(self.out_channel, channel))
 
        # 非关键字参数的特征是一个星号*加上参数名，比如*number，定义后，number可以接收任意数量的参数，并将它们储存在一个tuple中
        return nn.Sequential(*block)
 
    # 前向传播
    def forward(self, X):
        Y = self.relu(self.bn1(self.conv1(X)))
        Y = self.maxpool(Y)
        Y = self.conv5(self.conv4(self.conv3(self.conv2(Y))))
 
        if self.include_top:
            Y = self.avgpool(Y)
            Y = torch.flatten(Y, 1)
            # Y = self.dropout(Y)
            Y = self.fc(Y)
 
        return Y
 
 
# 构建ResNet模型
def resnet18(num_classes=10, include_top=True):
    return ResNet(BasicBlock, [2, 2, 2, 2], num_classes=num_classes, include_top=include_top)
def resnet34(num_classes=10, include_top=True):
    return ResNet(BasicBlock, [3 ,4 ,6 ,3], num_classes=num_classes, include_top=include_top)
def resnet50(num_classes=10, include_top=True):
    return ResNet(BottleNeck, [3, 4, 6, 3], num_classes=num_classes, include_top=include_top)
def resnet101(num_classes=10, include_top=True):
    return ResNet(BottleNeck, [3, 4, 23, 3], num_classes=num_classes, include_top=include_top)
# 同理定义 resnet34、resnet50、resnet101 


def main():
    # 创建输出目录
    num = 1
    while os.path.exists(f'./output{num}'):
        num += 1
    output_path = f'./output{num}'
    os.makedirs(output_path)

    models = [resnet18,resnet34,resnet50,resnet101]
    names = ['resnet18','resnet34','resnet50','resnet101']
    # 数据预处理
    train_transform = transforms.Compose([
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.RandomRotation(15),  # 随机旋转
        transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),  # 颜色抖动
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])

    val_transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])

    # 数据加载
    train_dataset = datasets.ImageFolder(root='./data_set_224/train', transform=train_transform)
    val_dataset = datasets.ImageFolder(root='./data_set_224/test', transform=val_transform)

    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=4)
    val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=4)

    for resnet, name in zip(models,names):
        start_time = time.time()
        #创建cvs文件
        with open(f'./output{num}/{name}_results.csv', "w", newline='') as csvfile:  
            writer = csv.writer(csvfile)  
            writer.writerow(["Epoch","Loss","Acc"])  

            # 模型初始化
            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
            model = resnet(num_classes=10).to(device)

            # 查看模型结构
            summary(model, (3, 224, 224))

            # 可视化模型计算图
            model.eval()  # 切换到评估模式，以确保不影响BatchNorm和Dropout层
            X = torch.rand(size=(1, 3, 224, 224)).to(device)  # 创建一个随机输入张量
            Y = model(X)
            vise = make_dot(Y, params=dict(model.named_parameters()))
            vise.format = 'png'
            vise.render(f'./output{num}/{name}_graph')  # 保存计算图到文件

            # 训练配置
            criterion = nn.CrossEntropyLoss()
            optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

            # optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=1e-4)

            scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)
            num_epochs = 100
            history = []

            # 训练循环
            best_acc = 0.0
            patience = 5
            counter = 0
            for epoch in range(num_epochs):
                # 训练阶段
                model.train()
                running_loss = 0.0
                for images, labels in train_loader:
                    images = images.to(device)
                    labels = labels.to(device)
                    optimizer.zero_grad()
                    outputs = model(images)
                    loss = criterion(outputs, labels)
                    loss.backward()
                    optimizer.step()
                    
                    running_loss += loss.item() * images.size(0)
                scheduler.step()
                epoch_loss = running_loss / len(train_dataset)
                
                # 验证阶段
                model.eval()
                correct = 0
                total = 0
                with torch.no_grad():
                    for images, labels in val_loader:
                        images = images.to(device)
                        labels = labels.to(device)
                        
                        outputs = model(images)
                        _, predicted = torch.max(outputs.data, 1)
                        total += labels.size(0)
                        correct += (predicted == labels).sum().item()
                
                epoch_acc = correct / total

                history.append(epoch_acc)
                # 写入csv文件
                writer.writerow([epoch+1,epoch_loss,epoch_acc])
                print(f'Epoch [{epoch+1}/{num_epochs}] Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}')
                # 保存最佳模型
                if epoch_acc > best_acc:
                    best_acc = epoch_acc
                    torch.save(model.state_dict(), f'./output{num}/{name}_best_model.pth')
                    counter = 0  # 重置计数器
                else:
                    counter += 1
                    if counter >= patience:
                        print(f'Early stopping at epoch {epoch+1}')
                        break    
            drewline(output_path, range(epoch+1), history,name)
            print(f'{name}Best Validation Accuracy: {best_acc:.4f}')
            end_time = time.time()
            epoch_duration = end_time - start_time
            print(f'Run time: {epoch_duration}')

def drewline(output_path, epoch, history,name):
    plt.figure(figsize=(8, 4.5))
    plt.plot(epoch, history, label=f'{name} Val Accuracy')
    plt.title(f'{name} Val Accuracy')
    plt.locator_params(axis='y', nbins=20)
    plt.savefig(os.path.join(output_path,f'{name}_derwline.png'))

if __name__ == '__main__':
    torch.multiprocessing.freeze_support()
    main()
